/**
 * @file main.cpp
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-10-26
 *
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 * This file can not be copied or distributed without the permission of Lynxi
 * Technologies Co., Ltd.
 *
 */
#include <lyn_api.h>
#include <fstream>
#include <iostream>
#include <nlohmann/json.hpp>
#include <vector>
#include "argsParser.h"
#include "imageDecoder.h"
#include "imageEncoder.h"
#include "ipeOcrParam.h"
#include "opencvWindow.hpp"
#include "osd_plugin.h"
#include "paddleocr_post_process.h"
#include "typeConv.hpp"
#include "util.hpp"

using namespace std;
using namespace ocr;

std::string GetBasePath(const std::string &filePath) {
  size_t lastSeparatorPos = filePath.find_last_of("/\\");
  if (lastSeparatorPos != std::string::npos) {
    return filePath.substr(0, lastSeparatorPos);
  }
  return "";
}

template <class ForwardIterator>
inline static size_t argmax(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::max_element(first, last));
}

void scoreToTextLine(const int16_t *outputData, std::string &res,
                     const std::vector<std::string> &mKeys, size_t model_h,
                     size_t model_w) {
  std::string strRes;
  // std::vector<float> scores;
  size_t lastIndex = 0;
  size_t maxIndex;
  // int16_t maxValue;
  std::vector<int> maxIndexVec;

  for (size_t i = 0; i < model_h; i++) {
    maxIndex = 0;
    // maxValue = -10000;

    maxIndex =
        int(argmax(outputData + i * model_w, outputData + (i + 1) * model_w));
    // maxValue = outputData[i * model_w + maxIndex];

    if (maxIndex > 0 && maxIndex < mKeys.size() &&
        (!(i > 0 && maxIndex == lastIndex))) {
      maxIndexVec.emplace_back(maxIndex);
      // scores.emplace_back(COMMON::half2float(maxValue));
      strRes.append(mKeys[maxIndex - 1]);
    }
    lastIndex = maxIndex;
  }

  res = strRes;
  // res.textLine.charScores = scores;

  return;
}

int main(int argc, char *argv[]) {
  // 1. 命令行解析
  if (ArgsParse::argsParser(argc, argv) == false) {
    return -1;
  }

  std::string strJpegPath = ArgsParse::argInputFilePath;

  // 2. 创建 context 与 stream
  lynContext_t ctx = nullptr;
  CHECK_ERR(lynCreateContext(&ctx, ArgsParse::argChipsVec[0]));
  CHECK_ERR(lynRegisterErrorHandler(StreamErrorHandler, nullptr));
  lynStream_t stream = nullptr;
  CHECK_ERR(lynCreateStream(&stream));

  std::string strJpegOutPath = strJpegPath;
  strJpegOutPath.insert(strJpegPath.rfind('.'), "_result");

  std::vector<std::string> mKeys;
  ifstream fileStream(ArgsParse::argKeysFilePath);
  if (!fileStream.is_open()) {
    throw invalid_argument(("paddleocr keys file read failed"));
  }

  std::string line;
  while (std::getline(fileStream, line)) {
    mKeys.emplace_back(line);
  }

  // 3. 初始化解码图片类
  ImageDecoder imageDec;
  imageDec.Init(strJpegPath);
  lynImageInfo_t imgInfo;
  imageDec.GetImageDecInfo(&imgInfo);

  // 4. 同步解码图片到 Device 侧
  lynCodecBuf_t decodeDevImg;
  decodeDevImg.size = imgInfo.output.predictBufSize;
  CHECK_ERR(lynMalloc((void **)&decodeDevImg.data, decodeDevImg.size));

  imageDec.DecodeImageToDevice(stream, decodeDevImg.data);

  // 5. 加载模型
  ModelInfo modelOcrDet;
  ModelInfo modelOcrCls;
  ModelInfo modelOcrRec;
  modelOcrDet.LoadModelByPath(ArgsParse::argDetModelPath.c_str());
  modelOcrCls.LoadModelByPath(ArgsParse::argClsModelPath.c_str());
  modelOcrRec.LoadModelByPath(ArgsParse::argRecModelPath.c_str());

  lynPixelFormat_t pixelFormat = LYN_PIX_FMT_NV12;

  // 6. 加载 Plugin, 并设置后处理参数
  lynPlugin_t postplugin;
  CHECK_ERR(lynPluginRegister(&postplugin, ArgsParse::argPostPluginPath.c_str()));
  lynPlugin_t osdplugin;
  CHECK_ERR(lynPluginRegister(&osdplugin, ArgsParse::argOsdPluginPath.c_str()));

  // 7. IPE 处理
  uint8_t *pDetIpeBufOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDetIpeBufOut, modelOcrDet.inputSize));
  IpeParamContain ipeParam(modelOcrDet.width, modelOcrDet.height);
  ipeParam.SetImgInfo(imgInfo.output.width, imgInfo.output.height, pixelFormat);
  ipeParam.CalcParam(stream, decodeDevImg.data, pDetIpeBufOut);

  // 8. 调用 APU 推理接口
  uint8_t *pDevApuOutData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevApuOutData, modelOcrDet.outputSize));
  CHECK_ERR(lynExecuteModelAsync(stream, modelOcrDet.model, pDetIpeBufOut,
                                 pDevApuOutData, modelOcrDet.batchSize));

  // detect后处理
  PaddleOcrDetectPostProcessInfo_t det_post_info;
  memset(&det_post_info, 0, sizeof(det_post_info));
  det_post_info.img_width = imgInfo.output.width;
  det_post_info.img_height = imgInfo.output.height;

  det_post_info.m_width = modelOcrDet.width;
  det_post_info.m_height = modelOcrDet.height;

  lynBoxesInfo *pDevBoxesInfo;
  CHECK_ERR(lynMalloc((void **)&pDevBoxesInfo, sizeof(lynBoxesInfo)));
  det_post_info.boxesInfo = pDevBoxesInfo;
  det_post_info.output_tensor = pDevApuOutData;
  det_post_info.image_data = decodeDevImg.data;

  // 抠图大小不超过原图，所以在这申请原图大小
  uint8_t *pDevCropImgData = nullptr;
  CHECK_ERR(lynMalloc((void **)&pDevCropImgData,
                      imgInfo.output.width * imgInfo.output.height * 3));
  det_post_info.crop_images_data = pDevCropImgData;

  cropImagesInfo *pDevCropImgInfo;
  CHECK_ERR(lynMalloc((void **)&pDevCropImgInfo, sizeof(cropImagesInfo)));
  det_post_info.cropsInfo = pDevCropImgInfo;

  CHECK_ERR(lynPluginRunAsync(stream, postplugin,
                              "lynPaddleOcrDetectPostProcess", &det_post_info,
                              sizeof(det_post_info)));

  lynBoxesInfo *pHostBoxesInfo = new lynBoxesInfo();
  cropImagesInfo hostCropImgInfo;

  lynMemcpyAsync(stream, pHostBoxesInfo, pDevBoxesInfo, sizeof(lynBoxesInfo),
                 ServerToClient);
  lynMemcpyAsync(stream, &hostCropImgInfo, pDevCropImgInfo,
                 sizeof(hostCropImgInfo), ServerToClient);
  lynSynchronizeStream(stream);

  // cls 和 rec推理
  int boxNum = pHostBoxesInfo->boxesNum;
  uint8_t *pClsRecIpeOut = nullptr;
  CHECK_ERR(lynMalloc((void **)&pClsRecIpeOut, modelOcrCls.inputSize * boxNum));

  uint8_t *pClsDevApuOutData = nullptr;
  CHECK_ERR(
      lynMalloc((void **)&pClsDevApuOutData, modelOcrCls.outputSize * boxNum));
  uint8_t *pClsHostApuOutData = new uint8_t[modelOcrCls.outputSize * boxNum];

  uint8_t *pRecDevApuOutData = nullptr;
  CHECK_ERR(
      lynMalloc((void **)&pRecDevApuOutData, modelOcrRec.outputSize * boxNum));
  lynMemset(pRecDevApuOutData, 0, modelOcrRec.outputSize * boxNum);
  uint8_t *pRecHostApuOutData = new uint8_t[modelOcrRec.outputSize * boxNum];

  // ipe处理
  IpeParamFixH ipeFixHParam(modelOcrCls.width, modelOcrCls.height);
  int offset = 0;
  for (int i = 0; i < boxNum; ++i) {
    const imageInfo &imgInfo = hostCropImgInfo.images[i];
    uint8_t *pIpeBufOut = pClsRecIpeOut + i * modelOcrCls.inputSize;
    ipeFixHParam.SetImgInfo(imgInfo.w, imgInfo.h, LYN_PIX_FMT_BGR24);
    ipeFixHParam.CalcParam(stream, pDevCropImgData + offset, pIpeBufOut);
    offset += (imgInfo.w * imgInfo.h * 3);
  }

  // cls推理
  for (int i = 0; i < boxNum; ++i) {
    CHECK_ERR(lynExecuteModelAsync(
        stream, modelOcrCls.model, pClsRecIpeOut + i * modelOcrCls.inputSize,
        pClsDevApuOutData + i * modelOcrCls.outputSize, modelOcrCls.batchSize));
  }

  // rec推理
  for (int i = 0; i < boxNum; ++i) {
    CHECK_ERR(lynExecuteModelAsync(
        stream, modelOcrRec.model, pClsRecIpeOut + i * modelOcrRec.inputSize,
        pRecDevApuOutData + i * modelOcrRec.outputSize, modelOcrRec.batchSize));
  }

  lynMemcpyAsync(stream, pClsHostApuOutData, pClsDevApuOutData,
                 modelOcrCls.outputSize * boxNum, ServerToClient);
  lynMemcpyAsync(stream, pRecHostApuOutData, pRecDevApuOutData,
                 modelOcrRec.outputSize * boxNum, ServerToClient);

  lynSynchronizeStream(stream);

  for (int i = 0; i < boxNum; i++) {
    // cls后处理
    uint16_t *pClsTmp =
        (uint16_t *)(pClsHostApuOutData + i * modelOcrCls.outputSize);
    std::string sCls = "0";
    if (*pClsTmp < *(pClsTmp + 1)) {
      sCls = "180";
    }

    // rec后处理
    std::string sRec;
    int16_t *pRecTmp =
        (int16_t *)(pRecHostApuOutData + i * modelOcrRec.outputSize);

    scoreToTextLine(pRecTmp, sRec, mKeys,
                    modelOcrRec.modelDesc->outputTensorAttrArray->dims[1],
                    modelOcrRec.modelDesc->outputTensorAttrArray->dims[2]);

    std::string strText = "text:" + sRec + " angle:" + sCls;
    memset(pHostBoxesInfo->boxes[i].label, 0, TEXT_MAX_LEN);
    strncpy(pHostBoxesInfo->boxes[i].label, strText.c_str(),
            std::min((int)strText.size(), TEXT_MAX_LEN));
  }

  lynFree(pClsRecIpeOut);
  lynFree(pClsDevApuOutData);
  lynFree(pRecDevApuOutData);
  delete[] pClsHostApuOutData;
  delete[] pRecHostApuOutData;

  lynMemcpy(pDevBoxesInfo, pHostBoxesInfo, sizeof(lynBoxesInfo),
            ClientToServer);

  for (int i = 0; i < boxNum; i++) {
    lynDrawBoxAndTextPara para;
    para.imgData = decodeDevImg.data;
    para.imgFmt = LYN_PIX_FMT_NV12;
    para.imgW = imgInfo.output.width;
    para.imgH = imgInfo.output.height;
    para.boxesInfo = pDevBoxesInfo;
    para.boxColor = DRAW_COLOR_BLUE;
    para.boxThick = DRAW_THICK_2;
    para.fontSize = FONT_SIZE_24;
    para.fontColor = DRAW_COLOR_BLUE;
    CHECK_ERR(lynPluginRunAsync(stream, osdplugin, "lynDrawBoxAndText", &para,
                                sizeof(para)));
  }

  lynSynchronizeStream(stream);

  // 9. 图片编码
  ImageEncoder encoder;
  encoder.EncodeImageAndSave(stream, strJpegOutPath, decodeDevImg, imgInfo);

  // 10. 销毁资源
  CHECK_ERR(lynDestroyStream(stream));
  CHECK_ERR(lynFree(decodeDevImg.data));
  CHECK_ERR(lynFree(pDetIpeBufOut));
  CHECK_ERR(lynFree(pDevApuOutData));
  CHECK_ERR(lynFree(pDevBoxesInfo));
  CHECK_ERR(lynFree(pDevCropImgData));
  CHECK_ERR(lynFree(pDevCropImgInfo));
  delete pHostBoxesInfo;
  modelOcrDet.UnLoadModel();
  modelOcrCls.UnLoadModel();
  modelOcrRec.UnLoadModel();
  if (ctx) {
    CHECK_ERR(lynDestroyContext(ctx));
  }

  return 0;
}
